#include <windows.h>
#include <iostream>
#include <tchar.h>

#include "XString.h"
#include "ApplicationsConfig.h"
#include "OutputConfig.h"
#include "DataCollectionTypes.h"
#include "RegistryHelper.h"

using namespace std;
using namespace platform;


/** global flag, the program will exit once set to TRUE */
static bool DONE = false;

/**
 * reads a number from standard input,
 * useful to read user's response when showing a menu
 */
unsigned int readOption(void) {
	unsigned int i;
	cin >> i;
	return i;
}

/** prints main program menu */
void printMainMenu(void) {
	cout << "+-----------+\n"
		 << "| MAIN MENU |\n"
		 << "+-----------+\n"
		 << "Please select operation:\n"
		 << "(1) Enable/disable monitoring\n"
		 << "(2) Configuration\n"
		 << "(3) Exit\n"
		 << "? ";
}

/** @return TRUE is such file exists, FALSE otherwise */
bool fileExists(const XString &dir, const XString &name) {
	XString fname(dir);
	fname += _T('\\');
	fname += name;
	DWORD dw = ::GetFileAttributes(fname.c_str());
	return (dw != 0xffffffff && (dw & FILE_ATTRIBUTE_DIRECTORY) == 0);
}


/**
 * This function validates that required files are already
 * copied to "system32" folder
 */
bool checkSystem32Files() {
	TCHAR buf[MAX_PATH];
	::GetEnvironmentVariable(_T("windir"), buf, MAX_PATH);
	XString sys32path(buf);
	sys32path += _T("\\system32");
	if (fileExists(sys32path, _T("abpa_loader.dll"))) {
		cout << "   abpa_loader.dll [OK]\n";
	} else {
		cout << "ERROR: file \"%windir%\\system32\\abpa_loader.dll\" could not be found, please copy the file first\n";
		return false;
	}

	if (fileExists(sys32path, _T("ABPAEngine.dll"))) {
		cout << "   ABPAEngine.dll [OK]\n";
	} else {
		cout << "ERROR: file \"%windir%\\system32\\ABPAEngine.dll\" could not be found, please copy the file first\n";
		return false;
	}

	if (fileExists(sys32path, _T("detoured.dll"))) {
		cout << "   detoured.dll [OK]\n";
	} else {
		cout << "ERROR: file \"%windir%\\system32\\detoured.dll\" could not be found, please copy the file first\n";
		return false;
	}
	return true;
}


/**
 * Shows list of options to enable/disable monitoring
 */
void printEnableDisableOptions(void) {
	cout << "+---------------------+\n"
		 << "| ENABLE/DISABLE MENU |\n"
		 << "+---------------------+\n"
		 << "Please select operation:\n"
		 << "(1) Enable ABPA\n"
		 << "(2) Temporary disable\n"
		 << "    (\"abpa_loader.dll\" will be still loaded in each process, but\n"
		 << "    will not add any hooks)\n"
		 << "(3) Disable and remove\n"
		 << "    (nothing will be loaded in new processes)\n"
		 << "(4) Main menu\n"
		 << "(5) Exit\n"
		 << "? ";
}


/**
 * This function is called when the user wants to set "monitoring"
 * registry flag to ON
 */
void monitoringFlagOn(void) {
	cout << "   enabling monitoring flag... ";
	config::ApplicationsConfig appconf;
	if (appconf.isMonitoringEnabled()) {
		cout << "[already enabled]\n";
		return;
	}
	appconf.enableMonitoring();
	appconf.save();
	cout << "[OK]\n";
}

/**
 * This function is called when the user wants to set "monitoring"
 * registry flag to OFF
 */
void monitoringFlagOff(void) {
	cout << "   disabling monitoring flag... ";
	config::ApplicationsConfig appconf;
	if (!appconf.isMonitoringEnabled()) {
		cout << "[already disabled]\n";
		return;
	}
	appconf.disableMonitoring();
	appconf.save();
	cout << "[OK]\n";
}

/**
 * Returns a list of DLLs registered in Windows to be injected to
 * each newly started process
 * SOFTWARE\Microsoft\Windows NT\CurrentVersion\Windows
 */
void getApiInitDlls(XStringList &dlls) {
	RegistryHelper::RegistryKey key(RegistryHelper::RK_LOCAL_MACHINE, _T("SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows"));
	XString dlls_string;
	key.getStringValue(_T("AppInit_DLLs"), dlls_string);
	dlls_string.split(_T(','), dlls);
}


/**
 * Updates the list of DLLs which will be injected to every new process
 */
void setApiInitDlls(XStringList &dlls) {
	RegistryHelper::RegistryKey key(RegistryHelper::RK_LOCAL_MACHINE, _T("SOFTWARE\\Microsoft\\Windows NT\\CurrentVersion\\Windows"));
	XString dlls_string;
	dlls.join(_T(','), dlls_string);
	key.setValue(_T("AppInit_DLLs"), dlls_string);
}


/**
 * Adds "abpa_loader.dll" to list of DLLs that will be
 * automatically injected to each Windows process
 */
void appInitAddToList(void) {
	cout << "   adding to \"AppInit_DLLs\" key... ";
	XStringList dlls;
	getApiInitDlls(dlls);

	// check whether is already in the list
	for (XStringList::iterator it = dlls.begin(); it != dlls.end(); ++it)
		if ((*it) == _T("abpa_loader.dll")) {
			cout << "[already contains \"abpa_loader.dll\"]\n";
			return;
		}

	// not in the list - should add (at the beginning)
	dlls.push_front(XString(_T("abpa_loader.dll")));
	setApiInitDlls(dlls);

	cout << "[OK]\n";
}

/**
 * Removes "abpa_loader.dll" from list of DLLs that will be
 * automatically injected to each Windows process
 */
void appInitRemoveFromList(void) {
	cout << "   removing from \"AppInit_DLLs\" key... ";
	XStringList dlls;
	getApiInitDlls(dlls);

	// check whether is already in the list
	bool found = false;
	XStringList::iterator it = dlls.begin();
	while (it != dlls.end()) {
		if ((*it) == _T("abpa_loader.dll")) {
			found = true;
			XStringList::iterator toerase = it;
			++it;
			dlls.erase(toerase);
		} else {
			++it;
		}
	}
	
	if (!found) {
		cout << "[\"abpa_loader.dll\" not found]\n";
		return;
	}

	// not in the list - should add (at the beginning)
	setApiInitDlls(dlls);
	cout << "[OK]\n";
}


/**
 * Shows and operates the menu to enable/disable ABPA monitoring
 */
void enableDisable(void) {
	cout << "\ntesting current settings...\n";
	if (!checkSystem32Files()) return;

	while (1) {
		printEnableDisableOptions();
		switch (readOption()) {
			case 1:
				monitoringFlagOn();
				appInitAddToList();
				cout << "done\n\n";
				break;
			case 2:
				monitoringFlagOff();
				cout << "done\n\n";
				break;
			case 3:
				monitoringFlagOff();
				appInitRemoveFromList();
				cout << "done\n\n";
				break;
			case 4:
				cout << endl;
				return;
			case 5:
				DONE = true;
				return;
		}
	}
}



/**
 * Prints list of selected applications (the actual setting can
 * be "use those only" or "use all except for those")
 */
void printSelectedApplications(config::ApplicationsConfig &conf) {
	const XStringList &apps = conf.getApplicationsList();
	if (apps.empty()) {
		cout << "      [empty list]\n";
		return;
	}

	for (XStringList::const_iterator it = apps.begin(); it != apps.end(); ++it) {
		cout << "      " << it->getANSIString() << endl;
	}
}


/**
 * Shows menu with options to edit applications' list
 */
void printSelectApplicationsMenu(void) {
	cout << "+-----------------------------+\n"
		 << "| CONFIGURATION: Applications |\n"
		 << "------------------------------+\n";

	cout << "Current settings:\n";
	config::ApplicationsConfig conf;
	switch (conf.getSelectionType()) {
		case config::ApplicationsConfig::AST_ALL:
			cout << "   Monitor ALL applications\n";
			break;
		case config::ApplicationsConfig::AST_ONLY_SELECTED:
			cout << "   Monitored only following applications:\n";
			printSelectedApplications(conf);
			break;
		case config::ApplicationsConfig::AST_EXCEPT_FOR_SELECTED:
			cout << "   Monitored all applications, EXCEPT for:\n";
			printSelectedApplications(conf);
			break;
	}

	cout << "Please select operation:\n"
		 << "(1) Change selection type\n"
		 << "(2) Edit applications' list\n"
		 << "(3) Test application\n"
		 << "(4) Back to Configuration menu\n"
		 << "(5) Exit\n"
		 << "? ";
}


/**
 * Asks the user to enter a name of some applications,
 * and answers whether this application should be monitored
 * with current settings.
 * NOTE: even when the answer is YES, nothing will be monitored
 * when the main "enable monitoring" flag is turned off.
 */
void testApplicationName(void) {
	config::ApplicationsConfig conf;
	
	while (1) {
		string name;
		cout << "Application name [\"q\" to return]: ";
		cin >> name;

		if (name == "q") {
			cout << endl;
			return;
		}

		if (conf.isApplicationEnabled(name)) 
			cout << "YES (this application will be monitored)\n";
		else
			cout << "NO (this application will not be monitored)\n";

	}
}


/**
 * Changes a way applications' list is regarded
 */
void changeAppSelectionType() {
	config::ApplicationsConfig conf;
	cout << "\nPlease choose selection type:\n"
		 << "(1) Monitor ALL applications";
	if (conf.getSelectionType() == config::ApplicationsConfig::AST_ALL) cout << " [current]";
	cout << "\n(2) Monitor only listed applications";
	if (conf.getSelectionType() == config::ApplicationsConfig::AST_ONLY_SELECTED) cout << " [current]";
	cout << "\n(3) Monitor all applications EXCEPT for listed";
	if (conf.getSelectionType() == config::ApplicationsConfig::AST_EXCEPT_FOR_SELECTED) cout << " [current]";
	cout << "\n(4) Cancel\n? ";

	switch (readOption()) {
		case 1:
			conf.setAppSelectionType(config::ApplicationsConfig::AST_ALL);
			break;
		case 2:
			conf.setAppSelectionType(config::ApplicationsConfig::AST_ONLY_SELECTED);
			break;
		case 3:
			conf.setAppSelectionType(config::ApplicationsConfig::AST_EXCEPT_FOR_SELECTED);
			break;
		case 4:
			return;
		default:
			cout << "<unknown option>\n";
			return;
	};

	conf.save();
}


/**
 * Menu for editing list of applications
 */
void editApplicationsList(void) {
	cout << "\n"
		 << "EDIT LIST OF APPLICATIONS\n"
		 << "(1) Add application\n"
		 << "(2) Remove application(s)\n"
		 << "? ";

	unsigned int opt = readOption();
	if (opt == 1) {
		// adding new application pattern
		string newApp;
		cout << "Please enter application name pattern: ";
		cin >> newApp;
		if (newApp.empty()) {
			cout << "<no patter was specified>\n";
			return;
		}

		config::ApplicationsConfig conf;
		conf.addApplication(newApp);
		return;
	}
	if (opt != 2) return;

	// removing application(s)
	config::ApplicationsConfig conf;
	do {
		const XStringList &apps = conf.getApplicationsList();
		cout << "\nREMOVING APPLICATION PATTERN\n";
		if (apps.empty()) {
			cout << "<empty list, nothing to remove>\n";
			return;
		}

		unsigned int curIndex = 1;
		for (XStringList::const_iterator it = apps.begin(); it != apps.end(); ++it) {
			cout << '(' << curIndex << ") Remove \"" << it->getANSIString() << "\"\n";
			curIndex ++;
		}
		unsigned int numApps = curIndex - 1;
		cout << '(' << curIndex << ") Remove ALL\n"
			 << '(' << (curIndex+1) << ") Done\n"
			 << "? ";
		opt = readOption();
		if (opt >= 1 && opt <= numApps)
			conf.removeApplication(opt - 1);
		else if (opt == (numApps+1))
			conf.clearAppList();
		else if (opt == (numApps + 2))
			return;

		if (opt > (numApps + 21)) {
			cout << "<unkown option>\n";
			return;
		}
	} while(1);
}


/**
 * Operates a menu for applications list management
 * (defines which applications should be monitored)
 */
void selectMonitoredApplications(void) {
	while (1) {
		printSelectApplicationsMenu();
		switch(readOption()) {
			case 1:
				changeAppSelectionType();
				break;
			case 2:
				editApplicationsList();
				break;
			case 3:
				testApplicationName();
				break;
			case 4:
				cout << endl;
				return;
			case 5:
				throw string("done");
		}
	}
}


/**
 * Configures location of ABPA output files
 */
void configOutput(void) {
	cout << "\nOUTPUT FILES\n"
		<< "Current location: ";
	config::OutputConfig conf;
	cout << conf.getLocation().getANSIString() << endl
		 << "(1) Change location\n"
		 << "(2) Cancel\n"
		 << "? ";

	if (readOption() == 1) {
		cout << "Please enter new location for output files\n";
		string newloc;
		cin >> newloc;
		if (newloc.empty()) {
			cout << "<empty location>\n";
			return;
		}
		conf.setLocation(newloc);
		conf.save();
	}
}


/**
 * Configures (enables/disables) different event types
 */
void configEventTypes(void) {
	config::DataCollectionTypes dtypes;
	while (1) {
		cout << "+-------------+\n"
			 << "| EVENT TYPES |\n"
			 << "+-------------+\n"
			 << "(1) File events             [" << (dtypes.shouldCollectFileEvents() ? "ON" : "OFF") << "]\n"
			 << "(2) Network events          [" << (dtypes.shouldCollectNetworkEvents() ? "ON" : "OFF") << "]\n"
			 << "(3) Virtual Memory events   [" << (dtypes.shouldCollectMemoryEvents() ? "ON" : "OFF") << "]\n"
			 << "(4) Thread events           [" << (dtypes.shouldCollectThreadEvents() ? "ON" : "OFF") << "]\n"
			 << "(5) Windows events          [" << (dtypes.shouldCollectWindowEvents() ? "ON" : "OFF") << "]\n"
			 << "(6) Synchronization events  [" << (dtypes.shouldCollectSyncronizationEvents() ? "ON" : "OFF") << "]\n"
			 << "(7) Done\n"
			 << "Please enter number of event to toggle, or '7' when finished\n"
			 << "? ";
		switch (readOption()) {
			case 1:
				if (dtypes.shouldCollectFileEvents())
					dtypes.disableFileEvents();
				else
					dtypes.enableFileEvents();
				break;
			case 2:
				if (dtypes.shouldCollectNetworkEvents())
					dtypes.disableNetworkEvents();
				else
					dtypes.enableNetworkEvents();
				break;
			case 3:
				if (dtypes.shouldCollectMemoryEvents())
					dtypes.disableMemoryEvents();
				else
					dtypes.enableMemoryEvents();
				break;
			case 4:
				if (dtypes.shouldCollectThreadEvents())
					dtypes.disableThreadEvents();
				else
					dtypes.enableThreadEvents();
				break;
			case 5:
				if (dtypes.shouldCollectWindowEvents())
					dtypes.disableWindowEvents();
				else
					dtypes.enableWindowEvents();
				break;
			case 6:
				if (dtypes.shouldCollectSyncronizationEvents())
					dtypes.disableSynchronizationEvents();
				else
					dtypes.enableSynchronizationEvents();
				break;
			default:
				return;
		}
		cout << endl;
	}
}


/**
 * Prints the main configuration menu
 */
void printConfigMenu(void) {
	cout << "+--------------------+\n"
		 << "| CONFIGURATION MENU |\n"
		 << "+--------------------+\n"
		 << "Please select operation:\n"
		 << "(1) Monitored applications\n"
		 << "(2) Output files\n"
		 << "(3) Event types\n"
		 << "(4) Main menu\n"
		 << "(5) Exit\n"
		 << "? ";
}


/**
 * Operates configuration menu
 */
void configure(void) {
	cout << endl;
	while (1) {
		printConfigMenu();
		switch(readOption()) {
			case 1:
				selectMonitoredApplications();
				break;
			case 2:
				configOutput();
				break;
			case 3:
				configEventTypes();
				break;
			case 4:
				cout << endl;
				return;
			case 5:
				DONE = true;
				return;
		}
	}
}



// MAIN
int main(int argc, char *argv[]) {
	cout << "=============================================\n"
		 << "Application Behavior Patterns Analyzer (ABPA)\n"
		 << "Configuration manager\n"
		 << "\n";

	try {
		while (!DONE) {
			printMainMenu();
			unsigned int opt = readOption();
			switch (opt) {
				case 1:
					enableDisable();
					break;
				case 2:
					configure();
					break;
				case 3:
					return 0;
			}
		}
	} catch(string&) {
	}

	return 0;
}
